import os
import numpy as np
import time
import torch

from classification import (Classification, predict_input)
from utils.utils_metrics import evaluteTop1_5



# ------------------------------------------------------#
#   test_annotation_path    测试图片路径和标签
# ------------------------------------------------------#
test_annotation_path = 'cls_test.txt'
# ------------------------------------------------------#
#   metrics_out_path        指标保存的文件夹
# ------------------------------------------------------#
metrics_out_path = "metrics_out"

class Eval_Classification(Classification):
    def detect_image(self, image):

        # ---------------------------------------------#
        #   对测试图片预处理
        # ---------------------------------------------#
        image_data = predict_input(image, self.input_shape)
        # ---------------------------------------------------------#
        #   添加batch_size维度
        # ---------------------------------------------------------#
        image_data = np.expand_dims(image_data, 0)

        with torch.no_grad():
            photo = torch.from_numpy(image_data).type(torch.FloatTensor)
            if self.cuda:
                photo = photo.cuda()
            # ---------------------------------------------------#
            #   图片传入网络进行预测
            # ---------------------------------------------------#
            # 计算每张图片的处理时间
            torch.cuda.synchronize()
            start = time.time()
            preds = torch.softmax(self.model(photo)[0], dim=-1)
            torch.cuda.synchronize()
            end = time.time()
            image_time = end - start
            preds = preds.cpu().numpy()
        return preds, image_time


if __name__ == "__main__":

    if not os.path.exists(metrics_out_path):
        os.makedirs(metrics_out_path)

    classfication = Eval_Classification()

    # 读取验证数据
    with open("./cls_valid.txt", "r") as f:
        val_lines = f.readlines()

    top1, top5, Recall, Precision, fps = evaluteTop1_5(classfication, val_lines, metrics_out_path)
    print("top-1 accuracy = %.2f%%" % (top1 * 100))
    print("top-5 accuracy = %.2f%%" % (top5 * 100))
    print("mean Recall = %.2f%%" % (np.mean(Recall) * 100))
    print("mean Precision = %.2f%%" % (np.mean(Precision) * 100))
    print("mean fps = %.2f" % fps)

    with open(os.path.join(metrics_out_path, "metrics.txt"), 'w') as f:
        f.write("top-1 accuracy = %.2f%%\n" % (top1 * 100))
        f.write("top-5 accuracy = %.2f%%\n" % (top5 * 100))
        f.write("mean Recall = %.2f%%\n" % (np.mean(Recall) * 100))
        f.write("mean Precision = %.2f%%\n" % (np.mean(Precision) * 100))
        f.write("mean fps = %.2f" % fps)
    print("Save accuracy, Recall, Precision out to " + os.path.join(metrics_out_path, "metrics.txt"))


